{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d4073749",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import SimpleDirectoryReader, VectorStoreIndex\n",
    "from llama_index.readers.file.base import (\n",
    "    DEFAULT_FILE_READER_CLS,\n",
    "    ImageReader,\n",
    ")\n",
    "from llama_index.response.notebook_utils import (\n",
    "    display_response,\n",
    "    display_image,\n",
    ")\n",
    "from llama_index.indices.query.query_transform.base import (\n",
    "    ImageOutputQueryTransform,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a4b74e87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: we add filename as metadata for all documents\n",
    "filename_fn = lambda filename: {\"file_name\": filename}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca801c8c",
   "metadata": {},
   "source": [
    "# Q&A over Receipt Images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80cce8e4",
   "metadata": {},
   "source": [
    "We first ingest our receipt images with the *custom* `image parser` and `metadata function` defined above.   \n",
    "This gives us `image documents` instead of only text documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dbc28f2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Could not find image processor class in the image processor config or the model config. Loading based on pattern matching with the model's feature extractor configuration.\n"
     ]
    }
   ],
   "source": [
    "receipt_reader = SimpleDirectoryReader(\n",
    "    input_dir=\"data/receipts\",\n",
    "    file_metadata=filename_fn,\n",
    ")\n",
    "receipt_documents = receipt_reader.load_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12fd6f45",
   "metadata": {},
   "source": [
    "We build a simple vector index as usual, but unlike before, our index holds images in addition to text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "629cab63",
   "metadata": {},
   "outputs": [],
   "source": [
    "receipts_index = VectorStoreIndex.from_documents(receipt_documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fef454f",
   "metadata": {},
   "source": [
    "We can now ask a question that prompts for response with both text and image.  \n",
    "We use a custom query transform `ImageOutputQueryTransform` to add instruction on how to display the image nicely in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7c078dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_engine import TransformQueryEngine\n",
    "\n",
    "\n",
    "query_engine = receipts_index.as_query_engine()\n",
    "query_engine = TransformQueryEngine(\n",
    "    query_engine, query_transform=ImageOutputQueryTransform(width=400)\n",
    ")\n",
    "receipts_response = query_engine.query(\n",
    "    \"When was the last time I went to McDonald's and how much did I spend?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05c180ae",
   "metadata": {},
   "source": [
    "We now have rich multimodal response with inline text and image!  \n",
    "\n",
    "The source nodes section gives additional details on retrieved data used for synthesizing the final response.  \n",
    "In this case, we can verify that the receipt for McDonald's is correctly retrieved. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "810ad2e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**`Final Response:`** <img src=\"data/receipts/1100-receipt.jpg\" width=\"400\" />\n",
       "\n",
       "The last time you went to McDonald's was on 03/10/2018 and you spent $26.15."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_response(receipts_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa834925",
   "metadata": {},
   "source": [
    "# Q & A over LlamaIndex Documentation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a82a54",
   "metadata": {},
   "source": [
    "We now demo the same for Q&A over LlamaIndex documentations.   \n",
    "This demo higlights the ability to synthesize multimodal output with a mixture of text and image documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d5f04295",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Could not find image processor class in the image processor config or the model config. Loading based on pattern matching with the model's feature extractor configuration.\n"
     ]
    }
   ],
   "source": [
    "llama_reader = SimpleDirectoryReader(\n",
    "    input_dir=\"./data/llama\",\n",
    "    file_metadata=filename_fn,\n",
    ")\n",
    "llama_documents = llama_reader.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "46db4191",
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_index = VectorStoreIndex.from_documents(llama_documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4a4cc090",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_engine import TransformQueryEngine\n",
    "\n",
    "\n",
    "query_engine = llama_index.as_query_engine(similarity_top_k=2)\n",
    "query_engine = TransformQueryEngine(\n",
    "    query_engine, query_transform=ImageOutputQueryTransform(width=400)\n",
    ")\n",
    "llama_response = query_engine.query(\n",
    "    \"Show an image to illustrate how tree index works and explain briefly.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "559624a6",
   "metadata": {},
   "source": [
    "By inspecting the 2 source nodes, we see relevant text and image describing the tree index are retrieved for synthesizing the final multimodal response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5c5721d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**`Final Response:`** <img src=\"data/llama/tree_index.png\" width=\"400\" />\n",
       "\n",
       "This image illustrates how a tree index works. It shows a hierarchical tree structure with a root node at the top and leaf nodes at the bottom. The nodes are connected by branches, which represent the relationships between the nodes. The tree index is a useful way to organize data in a hierarchical structure."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_response(llama_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbd7376e",
   "metadata": {},
   "source": [
    "We show another example asking about vector store index instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "92569825",
   "metadata": {},
   "outputs": [],
   "source": [
    "llama_response = query_engine.query(\n",
    "    \"Show an image to illustrate how vector store index works and explain briefly.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7cfdd68d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**`Final Response:`** <img src=\"data/llama/vector_store_index.png\" width=\"400\" />\n",
       "\n",
       "This image illustrates how vector store index works. It stores each Node and its corresponding embedding in a Vector Store. The Nodes are represented by circles and the embeddings are represented by arrows. The arrows point from the Node to the embedding, indicating that the Node is associated with the embedding."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_response(llama_response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f19cc672",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gpt_index",
   "language": "python",
   "name": "gpt_index"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
